package cn.xaut.ocr3.service.impl;


import cn.xaut.ocr3.dto.IdCard;
import cn.xaut.ocr3.service.RapidOcrPlusService;
import cn.xaut.ocr3.utils.OcrThreadPoolExecutor;
import cn.xaut.ocr3.utils.RapidOcrUtils;
import cn.xaut.ocr3.utils.RapidOcrUtilsAndExecutor;
import io.github.mymonstercat.Model;
import io.github.mymonstercat.ocr.InferenceEngine;
import io.github.mymonstercat.ocr.config.ParamConfig;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.rendering.PDFRenderer;
import org.apache.pdfbox.text.PDFTextStripper;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;

import javax.annotation.PostConstruct;
import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * @author dgl
 */
@Service
public class RapidOcrPlusServiceImpl implements RapidOcrPlusService {

    private final InferenceEngine engine;
    private final ParamConfig paramConfig;

    public RapidOcrPlusServiceImpl(InferenceEngine engine, ParamConfig paramConfig) {
        this.engine = engine;
        this.paramConfig = paramConfig;
    }

    String originPath;

    private static String extractInfo(String input, String regex) {
        Pattern pattern = Pattern.compile(regex);
        Matcher matcher = pattern.matcher(input);
        if (matcher.find()) {
            return matcher.group(1);
        }
        return "null";
    }

    @PostConstruct
    public void initOriginPath() {
        //获取临时文件路径
        originPath = System.getProperty("java.io.tmpdir") + File.separator + "ocrFile";
    }

    /**
     * 识别图片
     *
     * @param file
     * @return
     * @throws IOException
     */
    @Override
    public String recognizeImage(MultipartFile file) throws IOException, InterruptedException {

        List<String> list = new ArrayList<>();


        String fileName = file.getOriginalFilename();

        List<String> imageNames = Arrays.asList("jpg", "png", "jpeg");
        List<String> pdfNames = Arrays.asList("pdf", "PDF");
        String endFileName = fileName.substring(fileName.lastIndexOf(".") + 1).toLowerCase();
        //判断文件格式
        if (pdfNames.contains(endFileName)) {

            new File(originPath).mkdirs();
            File tempFile = new File(originPath + File.separator + file.getOriginalFilename());
            file.transferTo(tempFile);
            //将pdf文件转为图片
            PDDocument document = PDDocument.load(tempFile);
            PDFRenderer pdf = new PDFRenderer(document);
            int pageCount = document.getNumberOfPages();

            List<File> imageFiles = new ArrayList<>();

            //识别每一页
            for (int i = 0; i < pageCount; i++) {

                BufferedImage image = pdf.renderImageWithDPI(i, 300);
                File imageFile = new File(originPath + File.separator + i + ".jpg");
                ImageIO.write(image, "PNG", imageFile);
                imageFiles.add(imageFile);
            }

            OcrThreadPoolExecutor ocrExecutor = new OcrThreadPoolExecutor();
            List<Callable<String>> tasks = new ArrayList<>();

            for (File f : imageFiles) {
                tasks.add(() -> RapidOcrUtilsAndExecutor.recognizeImage(f, engine, paramConfig));
            }
            List<Future<String>> futures = ocrExecutor.invokeAll(tasks, 10, TimeUnit.MINUTES);
            //ocrExecutor.printPoolStatus();
            for (Future<String> future : futures) {
                try {
                    String result = future.get();
                    list.add(result);
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
            for (File f : imageFiles) {
                if (f.exists()) {
                    f.delete();
                }
            }
            if (tempFile.exists()) {
                tempFile.delete();
            }
            ocrExecutor.shutdown();
            document.close();

        }
        //照片
        else if (imageNames.contains(endFileName)) {

            new File(originPath).mkdirs();
            File tempFile = new File(originPath + File.separator + file.getOriginalFilename());
            file.transferTo(tempFile);
            //调用识别方法
            String result = RapidOcrUtils.recognizeImage(tempFile, engine, paramConfig);

            if (tempFile.exists()) {
                tempFile.delete();
            }
            list.add(result);
        }

        return list.toString().replaceAll("\\[|\\]", "");
    }

    @Override
    public IdCard recognizeID(MultipartFile frontFile, MultipartFile backFile) throws IOException, InterruptedException {

        String input = recognizeImage(frontFile);
        // 提取姓名
        String name = extractInfo(input, "姓名([^性]+)").trim();
        // 提取性别
        String gender = extractInfo(input, "性别([^民]+)").trim();
        // 提取民族
        String ethnicity = extractInfo(input, "民族([^出]+)").trim();
        // 提取出生日期
        String birthDate = extractInfo(input, "出生(\\d{4}年\\d{1,2}月\\d{1,2}日)([^住]+)").trim();
        // 提取住址
        String address = extractInfo(input, "住址([^公]+)");
        address = address.replaceAll(" ", "");
        address = address.replaceAll("\\s+", "");
        // 提取身份证号码
        String idNumber = extractInfo(input, "[^0-9Xx]([1-9]\\d{16}[\\dXx])(?![0-9Xx])");
        //识别反面
        String backInput = recognizeImage(backFile);
        String issue = extractInfo(backInput, "签发机关([^有效]+)");
        String validDate = extractInfo(backInput, "有效期限\\s*(\\d{4}[./年-]?\\d{2}[./月-]?\\d{2}[日]?\\s*-\\s*\\d{4}[./年-]?\\d{2}[./月-]?\\d{2}[日]?|[\\d]{4}[./年-]?[\\d]{2}[./月-]?[\\d]{2}[日]?\\s*至\\s*[\\d]{4}[./年-]?[\\d]{2}[./月-]?[\\d]{2}[日]?|\\d{8}-\\d{4}\\.\\d{2}\\.\\d{2})");
        IdCard  idcard = new IdCard();
        idcard.setName("null".equals(name) ? null : name);
        idcard.setSex("null".equals(gender) ? null : gender);
        idcard.setNation("null".equals(ethnicity) ? null : ethnicity);
        idcard.setBirth("null".equals(birthDate) ? null : birthDate);
        idcard.setAddress("null".equals(address) ? null : address);
        idcard.setNumber("null".equals(idNumber) ? null : idNumber);
        idcard.setIssue("null".equals(issue) ? null : issue);
        idcard.setValidDate("null".equals(validDate) ? null : validDate);
        return idcard;
    }

    @Override
    public String recognizePDF(MultipartFile file) throws IOException {
        if (file != null) {
            new File(originPath).mkdirs();
            File tempFile = new File(originPath + File.separator + file.getOriginalFilename());
            file.transferTo(tempFile);
            PDDocument document = PDDocument.load(tempFile);
            PDFTextStripper stripper = new PDFTextStripper();
            stripper.setSortByPosition(true);
            String text = stripper.getText(document);
            document.close();
            if (tempFile.exists()) {
                tempFile.delete();
            }
            return text;
        }
        return "";
    }

    @Override
    public Object recognizeBankCard(MultipartFile file) throws IOException, InterruptedException {
        String card = recognizeImage(file);
        card = card.replaceAll(" ", "");
        System.out.println(card);
        return card;
    }

}
